package com.zt.chatgpt.openai.service.rule.impl;

import com.zt.chatgpt.openai.annotation.LogicStrategy;
import com.zt.chatgpt.openai.model.aggregates.ChatProcessAggregate;
import com.zt.chatgpt.openai.model.aggregates.UserAccountQuotaEntity;
import com.zt.chatgpt.openai.model.entity.RuleLogicEntity;
import com.zt.chatgpt.openai.model.valobj.LogicCheckTypeVO;
import com.zt.chatgpt.openai.service.rule.ILogicFilter;
import com.zt.chatgpt.openai.service.rule.factory.DefaultLogicFactory;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * @author: zt
 * @Description: 账户可用模型过滤
 * @create: 2024-05-22 20:23
 */
@Slf4j
@Component
@LogicStrategy(logicMode = DefaultLogicFactory.LogicModel.MODEL_TYPE)
public class ModelTypeFilter implements ILogicFilter<UserAccountQuotaEntity> {

    @Override
    public RuleLogicEntity<ChatProcessAggregate> filter(ChatProcessAggregate chatProcessAggregate, UserAccountQuotaEntity data) throws Exception {
        // 1. 获取用户请求模型和用户可用模型列表
        List<String> allowModelTypes = data.getAllowModelTypes();
        String modelType = chatProcessAggregate.getModel();

        // 2. 判断用户请求模型是否在用户可用模型列表中
        if (allowModelTypes.contains(modelType)) {
            return RuleLogicEntity.<ChatProcessAggregate>builder()
                    .type(LogicCheckTypeVO.SUCCESS)
                    .info("当前模型可用")
                    .data(chatProcessAggregate)
                    .build();
        }

        // 3. 模型不可用
        return RuleLogicEntity.<ChatProcessAggregate>builder()
                .type(LogicCheckTypeVO.REFUSE)
                .info("当前账户不支持使用 " + modelType + " 模型！可以联系客服升级账户。")
                .data(chatProcessAggregate)
                .build();
    }
}
